import numpy as np
# import pykalman
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import matplotlib as mpl
import scipy.stats as stats
import ekf_testv6 as ekf
import copy

data_CO2x4 = np.genfromtxt(open("CESM2-4xCO2.csv", "rb"),dtype=float, delimiter=',') 
temps = data_CO2x4[1,1:]
N = 30
moving_aves=np.zeros(len(temps))
std_aves=np.zeros(len(temps))
cN=int(np.ceil(N/2))
fN=int(np.floor(N/2))
for i in range((fN),(len(temps)-cN+1)):
    lasta=i+cN;firsta=i-fN
    moving_aves[i] = np.mean(temps[(firsta):lasta])
    std_aves[i] = np.std(temps[(firsta):lasta])

ax = plt.gca()
plt.rcParams['figure.figsize'] = (5,4)
plt.subplots_adjust(left=0.2, right=0.8)
updN=10
ax.set_xlabel('Model Year')
ax.set_ylabel('Temperature (K)')
ax.set_title('4xCO2 Experiment in CESM2')
ax.tick_params( direction = 'in',bottom=True, top=True, left=True, right=True )
mn,mx=287,295
plt.ylim(mn,mx)
plt.yticks(np.arange(mn,mx+1,1))


ax.plot(np.arange(475,550),temps[75:150],'o',label='Simulated CESM2 GMST',markersize=2,color=ekf.colorgrey)
ax.plot(np.arange(475,550),moving_aves[75:150],label='running 30-year average',markersize=2,color="mediumseagreen")
ax.plot(np.arange(475,550),moving_aves[75:150],'--',color='yellow',linewidth=0.8)
ax.fill_between(np.arange(475,550),moving_aves[75:150]+2*std_aves[75:150]/np.sqrt(30),moving_aves[75:150]-2*std_aves[75:150]/np.sqrt(30),label='30-year std error', color=ekf.colorstate)
handles0, labels = ax.get_legend_handles_labels()
dotted_line1 = plt.Line2D([], [],  linestyle="--", color="yellow", linewidth=0.8)
dotted_line2 = plt.Line2D([], [],  linestyle="-", color="mediumseagreen")
handles0[1] = (dotted_line2, dotted_line1)
ax.legend(handles0,labels)
ax1b = ax.twinx()
ax1b.set_ylim(mn-273.15, mx-273.15)
ax1b.set_ylabel('Temperature (°C)')


fig, ax = plt.subplots(figsize=(7,5))
plt.subplots_adjust(left=0.2, right=0.8)
ax.set_xlabel('Year')
ax.set_ylabel('Temperature (K)')
xhatorig = copy.deepcopy(ekf.xhat[:,0])
ax.plot(ekf.dates,xhatorig,'-',label='posterior GMST EBM-KF state \n$\hat{T_t}$ with time-varying TSI',color=ekf.colorekf)
ekf.tsi[:] = ekf.sw_in
this_xhat=ekf.ekf_run(ekf.observ,ekf.n_iters)
ax.plot(ekf.dates,this_xhat[:,0],'-',label='posterior GMST EBM-KF state \n with TSI=340.2W/m2')
ax.legend()
ax1b = ax.twinx()
ax1b.set_ylim(mn-273.15, mx-273.15)
ax1b.set_ylabel('Temperature (°C)')

plt.show()


